from collections import defaultdict, Counter
import math

class NaiveBayesClassifier:
    def __init__(self):
        self.class_prior = {}  # 类先验概率
        self.feature_prob = defaultdict(Counter)  # 类条件概率：{类别: {特征: 计数}}
        self.all_features = set()  # 所有特征集合
        self.classes = set()  # 所有类别集合

    # 训练模型
    def fit(self, X, y):
        # X: 特征列表（每个元素是关键词列表），y: 类别列表
        total_samples = len(y)
        self.classes = set(y)
        self.all_features = set([word for doc in X for word in doc])
        feature_count = len(self.all_features)

        # 计算先验概率 P(Y=c_k)
        class_count = Counter(y)
        for cls in self.classes:
            self.class_prior[cls] = class_count[cls] / total_samples

        # 计算类条件概率 P(X=x|Y=c_k)（拉普拉斯平滑）
        for doc, cls in zip(X, y):
            for word in doc:
                self.feature_prob[cls][word] += 1

        # 平滑处理：避免概率为0
        for cls in self.classes:
            total_word = sum(self.feature_prob[cls].values())
            for word in self.all_features:
                self.feature_prob[cls][word] = (self.feature_prob[cls][word] + 1) / (total_word + feature_count)

    # 预测单个样本
    def predict(self, x):
        max_prob = -float('inf')
        best_cls = None
        for cls in self.classes:
            # 计算后验概率（取对数避免下溢）
            log_prob = math.log(self.class_prior[cls])
            for word in x:
                if word in self.all_features:
                    log_prob += math.log(self.feature_prob[cls][word])
            if log_prob > max_prob:
                max_prob = log_prob
                best_cls = cls
        return best_cls

# 加载数据集
def load_dataset(file_path):
    X = []
    y = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip().split()
            if line:
                y.append(line[0])  # 第一列为类别
                X.append(line[1:])  # 后续为特征关键词
    return X, y

# 测试
if __name__ == "__main__":
    # 加载数据
    dataset_path = "datasets/20newsgroups_simplified.txt"
    X_train, y_train = load_dataset(dataset_path)

    # 训练模型
    nb = NaiveBayesClassifier()
    nb.fit(X_train, y_train)

    # 测试样本
    test_samples = [
        ["篮球", "球员", "决赛"],
        ["手机", "处理器", "续航"],
        ["网球", "球拍", "发球"],
        ["电脑", "显卡", "内存"]
    ]

    # 预测并输出结果
    print("朴素贝叶斯算法预测结果：")
    for sample in test_samples:
        pred = nb.predict(sample)
        print(f"样本{sample} → 预测类别：{pred}")